from collections import deque
import os
import database
import modelopr
import preprocess
from threading import Thread


def train_model(model_name, dataset_id, num_epochs, batch_size, learning_rate, optimizer, train_ratio, val_ratio,
                test_ratio):
    dataset = preprocess.get_dataset(dataset_id, train_ratio, val_ratio)

    def task():
        test_acc = modelopr.train(model_name=model_name, dataset=dataset, lr=learning_rate, num_epochs=num_epochs,
                                  batch_size=batch_size, optimizer=optimizer).item()

        database.insert_model(
            [model_name, dataset_id, train_ratio, val_ratio, num_epochs, batch_size, learning_rate, optimizer,
             test_acc])

    train_thread = Thread(target=task)
    train_thread.start()


def get_train_process(model_name="lenet") -> tuple[list[modelopr.TrainItem], bool]:
    q: deque = modelopr.train_queues[model_name]
    if deque is None:
        raise ValueError("No progress to provide.")
    length = len(q)
    ret = []
    for i in range(length):
        ret.append(q.popleft())
    finished = len(ret) > 0 and ret[-1].now_epoch == modelopr.total_epochs[model_name]
    print(modelopr.total_epochs[model_name])
    modelopr.train_queues[model_name].clear()
    if finished:
        modelopr.train_queues[model_name] = None
        modelopr.total_epochs[model_name] = None
    return ret, finished


def test_img(model_name, img_file):
    os.makedirs("tmp", exist_ok=True)
    file_path = os.path.join("tmp", "test.png")
    img_file.save(file_path)
    return modelopr.test(model_name, input_img=file_path)


if __name__ == "__main__":
    train_model("LeNet", 1, num_epochs=50, learning_rate=0.01, optimizer="adam", train_ratio=0.7, val_ratio=0.1,
                test_ratio=0.2, batch_size=32)
